"""
根据profile的结果进行模型的拆分，并验证和之前没有拆分之前的结果比对
"""
import torch
import torch.fx as fx
import graphviz
from parse_config import Config
import shutil
import os
import find_greatest_split

# --- 可视化函数 ---
def visualize_fx_graph(fx_graph, output_file, dpi=300):
    """
    将 fx.Graph 对象转换为 Graphviz 格式并保存为 PNG 文件。
    参数：
    fx_graph    : torch.fx.Graph 对象
    output_file : 输出的文件名前缀（如 'graph_before'）
    dpi         : 每英寸点数，控制转换英寸到像素时的比例
    """
    dot = graphviz.Digraph()
    dot.attr(rankdir='TB', dpi=str(dpi))
    for node in fx_graph.nodes:
        dot.node(node.name, str(node))
        for input_node in node.all_input_nodes:
            dot.edge(input_node.name, node.name)
    dot.render(output_file, format='png', cleanup=True)
    print(f"Graph saved as {output_file}.png")

def split_model(model_path,node_name):
    # 加载原始模型
    loaded_traced = torch.load(model_path)

    # 获取计算图
    graph = loaded_traced.graph

    # 选择要分割的节点
    target_node_name =node_name

    # 1. 找到指定的 node
    target_node = None
    for node in graph.nodes:
        if node.name == target_node_name:
            target_node = node
            break
    if target_node is None:
        raise ValueError(f"Node with name '{target_node_name}' not found in the graph.")

    # 2. 分割计算图
    # 创建第一个图 (target_node 之前的部分)
    graph_before = fx.Graph()
    node_mapping = {}  # 用于映射原始图中的 node 到新图中的 node


    # 复制 target_node 之前的部分
    for node in graph.nodes:
        new_node = graph_before.node_copy(node, lambda n: node_mapping[n])
        node_mapping[node] = new_node
        if node == target_node:
            break

    # 在 graph_before 的末尾添加输出节点
    with graph_before.inserting_after(list(graph_before.nodes)[-1]):
        graph_before.output(node_mapping[target_node])

    # 创建第二个图 (target_node 之后的部分)
    graph_after = fx.Graph()
    node_mapping = {}  # 重新初始化 node_mapping

    # 在 graph_after 的开头添加输入节点 (确保名称是 'input_node')
    input_node = graph_after.placeholder('input_node')

    flag=False
    # 复制 target_node 之后的部分
    flag = False
    for node in graph.nodes:
        if node == target_node:
            flag = True
            continue  # 跳过 target_node
        if not flag:
            continue
        new_node = graph_after.node_copy(
            node,
            lambda n: input_node if n == target_node or n not in node_mapping else node_mapping[n]
        )
        node_mapping[node] = new_node

    # 3. 创建新的 fx.GraphModule
    module_before = fx.GraphModule(loaded_traced, graph_before)
    module_after = fx.GraphModule(loaded_traced, graph_after)
    
    torch.save(module_before, 'module_before.pth')
    torch.save(module_after, 'module_after.pth')
        
    visualize_fx_graph(module_before.graph, 'graph_before')
    visualize_fx_graph(module_after.graph, 'graph_after')

    # 4. 运行分割后的图
    input_data = torch.randn(1, 3, 224, 224)

    # 运行第一个图
    output_before = module_before(input_data)

    # 运行第二个图
    output_after = module_after(input_node=output_before)

    # 5. 验证结果
    # 运行原始模型
    original_output = loaded_traced(input_data)

    # 检查分割后的结果是否与原始结果一致
    assert torch.allclose(output_after, original_output, atol=1e-6), "Outputs do not match!"
    print("Outputs match! The graph was split correctly.")
    
def split_model_segmentation(model_path,start_node_name,end_node_name,i,model_save_dir,graph_save_dir):
    """
    将模型从start_node拆分到end_nod，需要包含末节点

    Args:
        model_path (_type_): 模型输入路径
        start_node (_type_): 开始节点
        end_node (_type_): 终止节点
        i (_type_): 当前是第几个图
    """
    loaded_traced = torch.load(model_path)
    graph = loaded_traced.graph
    graph_result = fx.Graph()
    node_mapping = {}
    
    flag=False
    
    start_node = None
    for node in graph.nodes:
        if node.name == start_node_name:
            start_node = node
            break
    if start_node is None:
        raise ValueError(f"Node with name '{start_node_name}' not found in the graph.")
    
    end_node = None
    for node in graph.nodes:
        if node.name == end_node_name:
            end_node = node
            break
    if end_node is None:
        raise ValueError(f"Node with name '{end_node_name}' not found in the graph.")
    
    if i == 0:
        for node in graph.nodes:
            new_node = graph_result.node_copy(node, lambda n: node_mapping[n])
            node_mapping[node] = new_node
            if node == end_node:
                break

        # 在 graph_before 的末尾添加输出节点
        with graph_result.inserting_after(list(graph_result.nodes)[-1]):
            graph_result.output(node_mapping[end_node])
            
        module_result = fx.GraphModule(loaded_traced, graph_result)
        visualize_fx_graph(module_result.graph, f'{graph_save_dir}/graph_result{i}')
        module_result = fx.GraphModule(loaded_traced, graph_result)
        torch.save(module_result, os.path.join(model_save_dir,f'module_result{i}.pth'))
        return
    
    graph_tmp=fx.Graph()
    for node in graph.nodes:
        new_node = graph_tmp.node_copy(node, lambda n: node_mapping[n])
        node_mapping[node] = new_node
        if node == start_node:
            break
        
    if i != 0:
        input_node = graph_result.placeholder(f'input_node{i}')

        
    for node in graph.nodes:
        if node == start_node:
            flag = True
            continue  # 跳过 target_node
        if not flag:
            continue
        new_node = graph_result.node_copy(
            node,
            lambda n: input_node if n == start_node and i != 0 or n not in node_mapping else node_mapping[n]
        )
        node_mapping[node] = new_node
        if node == end_node:
            break
        
    with graph_result.inserting_after(list(graph_result.nodes)[-1]):
        graph_result.output(node_mapping[end_node])
        
    module_result = fx.GraphModule(loaded_traced, graph_result)
    visualize_fx_graph(module_result.graph, f'{graph_save_dir}/graph_result{i}')
    module_result = fx.GraphModule(loaded_traced, graph_result)
    torch.save(module_result, os.path.join(model_save_dir,f'module_result{i}.pth'))
    
def generate_model_and_graph():
    config_path='./config.json'
    config=Config(config_path)
    shutil.rmtree(config.config['split_module_save_dir'], ignore_errors=True)
    os.makedirs(config.config['split_module_save_dir'],exist_ok=True)
    shutil.rmtree(config.config['split_graph_save_dir'],ignore_errors=True)
    os.makedirs(config.config['split_graph_save_dir'],exist_ok=True)
    model_path = './results/formal_forward_graph.pth'
    cpu_nodes, node_list=find_greatest_split.main()
    node_list.append('output')
    cur_state=False
    cur_node=0
    cnt=0
    for i in range(len(node_list)):
        if i==0:
            if i in cpu_nodes:
                cur_state=True  
            continue
        if i == len(node_list)-1:
            split_model_segmentation(model_path,node_list[cur_node],node_list[i],cnt,config.config['split_module_save_dir'],config.config['split_graph_save_dir'])
            cnt+=1
        if i in cpu_nodes and not cur_state:
            cur_state=True
            split_model_segmentation(model_path,node_list[cur_node],node_list[i-1],cnt,config.config['split_module_save_dir'],config.config['split_graph_save_dir'])
            cnt+=1
            cur_node=i-1
            continue
        if i not in cpu_nodes and cur_state:
            cur_state=False
            split_model_segmentation(model_path,node_list[cur_node],node_list[i-1],cnt,config.config['split_module_save_dir'],config.config['split_graph_save_dir'])
            cnt+=1
            cur_node=i-1
            continue    
        
    
def main():
    generate_model_and_graph()
            

if __name__ == '__main__':
    main()
    
